package huobi

import (
	"bytes"
	"compress/gzip"
	"encoding/json"
	"fmt"
	"github.com/morgine/log"
	"io"
	"io/ioutil"
	"market/pkg/ws"
	"strings"
)

const (
	DepthLevel5  Level = 5
	DepthLevel10 Level = 10
	DepthLevel20 Level = 20
)

type Level int

func SubscribeDepth(symbols []string, level Level, handle func(data *DepthData)) (cancel func() error) {
	var subMsgs [][]byte
	var unsubMsgs [][]byte
	id := 1000
	//  "market.btcusdt.mbp.refresh.20"
	for _, symbol := range symbols {
		topic := fmt.Sprintf("market.%s.mbp.refresh.%d", symbol, level)
		subMsgs = append(subMsgs, []byte(fmt.Sprintf("{\"sub\": \"%s\",\"id\": \"%d\" }", topic, id)))
		unsubMsgs = append(unsubMsgs, []byte(fmt.Sprintf("{\"unsub\": \"%s\",\"id\": \"%d\" }", topic, id)))
	}
	return subscribeWs("wss://api.huobi.pro/ws", subMsgs, unsubMsgs, func(msg *Message) {
		depth, err := msg.DepthData()
		if err != nil {
			log.Error.Println(err)
		} else {
			handle(depth)
		}
	})
}

type Message struct {
	Ping      int64           `json:"ping"`
	Ch        string          `json:"ch"`
	Timestamp int64           `json:"ts"`
	Tick      json.RawMessage `json:"tick"`
}

func (m *Message) DepthData() (data *DepthData, err error) {
	data = &DepthData{}
	err = json.Unmarshal(m.Tick, data)
	if err != nil {
		return
	}
	// market.btcusdt.mbp.refresh.20
	ch := strings.TrimPrefix(m.Ch, "market.")
	data.Symbol = ch[:strings.Index(ch, ".")]
	data.Timestamp = m.Timestamp
	return
}

type DepthData struct {
	Timestamp int64         `json:"timestamp"`
	Symbol    string        `json:"symbol"`
	Asks      []DepthOrders `json:"asks"`
	Bids      []DepthOrders `json:"bids"`
}

type DepthOrders []float64

func (os DepthOrders) Price() float64 {
	return os[0]
}

func (os DepthOrders) Size() float64 {
	return os[1]
}

func subscribeWs(url string, connectMsgs, closeMsgs [][]byte, handler func(msg *Message)) (cancelFunc func() error) {
	var wsClient *ws.Client
	wsClient = ws.NewClient(url, func(w io.Writer) {
		for _, msg := range connectMsgs {
			_, err := w.Write(msg)
			if err != nil {
				log.Error.Println(err)
			}
		}
	}, func(message []byte) {
		var err error
		message, err = gzipDecompress(message)
		if err != nil {
			log.Error.Println(err)
		} else {
			msg := &Message{}
			err = json.Unmarshal(message, msg)
			if err != nil {
				log.Error.Println(err)
			} else {
				if msg.Ping != 0 {
					pong, _ := json.Marshal(struct {
						Pong int64 `json:"pong"`
					}{
						Pong: msg.Ping,
					})
					_, err = wsClient.Write(pong)
					if err != nil {
						log.Error.Println(err)
					}
				} else {
					if len(msg.Tick) == 0 {
						log.Info.Printf("huobi: %s\n", message)
					} else {
						handler(msg)
					}
				}
			}
		}
	})
	wsClient.Run()
	return func() error {
		for _, msg := range closeMsgs {
			wsClient.Write(msg)
		}
		return wsClient.Close()
	}
}

func gzipDecompress(input []byte) ([]byte, error) {
	reader, err := gzip.NewReader(bytes.NewReader(input))
	if err != nil {
		return nil, err
	}
	defer reader.Close()
	return ioutil.ReadAll(reader)
}
